#include <QtWidgets>
#include "mainwindow.h"
#include "ui_mainwindow.h"
#include "canvas2d.h"
#include "canvas3d.h"
#include "rectannotationitem.h"


MainWindow::MainWindow(QWidget *parent)
    : QMainWindow(parent)
    , ui(new Ui::MainWindow),
    labelManager(this),
    annoContainer(this),
    fileManager(this)
{
    ui->setupUi(this);

    fileRelatedActions << ui->actionSave
                       << ui->actionSave_As
                       << ui->actionClose
                       << ui->actionLoad
                       << ui->actionPrevious_Image
                       << ui->actionNext_Image
                       << ui->actionZoom_In
                       << ui->actionZoom_Out
                       << ui->actionFit_Window;

    canvas2d = new Canvas2D(&labelManager, &annoContainer, ui->scrollArea);
    canvas2d->setVisible(true); canvas2d->setEnabled(true);

    canvas3d = new Canvas3D(&labelManager, &annoContainer, ui->scrollArea);
    canvas3d->setVisible(false); canvas3d->setEnabled(false);

    curCanvas = canvas2d;
    ui->scrollArea->setAlignment(Qt::AlignCenter);
    ui->scrollArea->setWidget(curCanvas);

    zoomWidget = new ZoomWidget(this);
//    zoom = new QWidgetAction(this);
//    zoom->setDefaultWidget(zoomWidget);
//    ui->toolBar->addAction(zoom);
    zoomWidget->setEnabled(false);
    ui->toolBar->addWidget(zoomWidget);

//    _setupToolBarAndStatusBar();
    _setupLabelManager();
    _setupAnnotationContainer();
    _setupFileManager();

//    canvas2d->changeTask(DETECTION);        // DEFAULT TASK
    unableFileActions();
}

MainWindow::~MainWindow()
{
    delete ui;
    delete zoomWidget;
}

void MainWindow::_setupLabelManager()
{
    ui->labelListWidget->setSortingEnabled(true);

    // label的右键菜单: change color & delete
    ui->labelListWidget->setContextMenuPolicy(Qt::CustomContextMenu);
    connect(ui->labelListWidget, &QListWidget::customContextMenuRequested,
            this, &MainWindow::provideLabelContextMenu);

    // label的visible改变
    connect(ui->labelListWidget, &QListWidget::itemChanged,
            [this](QListWidgetItem *item){
                if (item->checkState()==Qt::Checked){
                    labelManager.setVisible(item->text(),true);
                }else{
                    labelManager.setVisible(item->text(),false);
                }
            });

    // label changed -> canvas update
    connect(&labelManager, &LabelManager::labelChanged, this, &MainWindow::canvasUpdate);

    // label changed -> ui list changed
    connect(&labelManager, &LabelManager::labelAdded,
            ui->labelListWidget, &CustomListWidget::addCustomItem);
    connect(&labelManager, &LabelManager::labelRemoved,
            ui->labelListWidget, &CustomListWidget::removeCustomItem);
    connect(&labelManager, &LabelManager::colorChanged,
            ui->labelListWidget, &CustomListWidget::changeIconColor);
    connect(&labelManager, &LabelManager::visibelChanged,
            ui->labelListWidget, &CustomListWidget::changeCheckState);
    connect(&labelManager, &LabelManager::allCleared,
            ui->labelListWidget, &QListWidget::clear);

    // 执行undo的时候可能会造成已经被删除的label又重新出现的情况
    connect(&annoContainer, &AnnotationContainer::labelGiveBack, this, &MainWindow::newLabelRequest);

    // 与新建label相关的部件
    connect(ui->labelButtonAdd, &QPushButton::clicked, [this](){
        newLabelRequest(ui->labelLineEdit->text());
        ui->labelLineEdit->setText("");
    });
    ui->labelLineEdit->installEventFilter(this); // 详见 MainWindow::eventFilter
}

bool MainWindow::eventFilter(QObject *watched, QEvent *event)
{
    if (watched == ui->labelLineEdit) {
        if (event->type() == QEvent::KeyPress) {
            QKeyEvent *keyEvent = static_cast<QKeyEvent*>(event);
            if (keyEvent->key() == Qt::Key_Return || keyEvent->key() == Qt::Key_Enter) {
                newLabelRequest(ui->labelLineEdit->text());
                ui->labelLineEdit->setText("");
                return true;    // 不继续传播回车事件是因为enter也可用于表示一个分割标注已经绘制完毕
            }
        }
    }
    return false;
}

void MainWindow::_setupAnnotationContainer()
{
    ui->annoListWidget->setSortingEnabled(false);

    // move up 和 move bottom 用于改变分割标注的覆盖情况，所以只在分割且选中时是enabled
    ui->annotationButtonMoveUp->setEnabled(false);
    ui->annotationButtonMoveDown->setEnabled(false);
    connect(ui->annotationButtonMoveUp, &QPushButton::clicked, [this](){
        int idx = annoContainer.getSelectedIdx();
        if (idx!=-1 && idx!=0){
            annoContainer.swap(idx-1);
        }
    });
    connect(ui->annotationButtonMoveDown, &QPushButton::clicked, [this](){
        int idx = annoContainer.getSelectedIdx();
        if (idx!=-1 && idx!=annoContainer.length()-1){
            annoContainer.swap(idx);
        }
    });

    // 选中一个anno的item来进入选择编辑模式
    connect(ui->annoListWidget, &QListWidget::itemSelectionChanged, [this](){
        auto items = ui->annoListWidget->selectedItems();
        if (items.length()==0){ // 回到DRAW模式
            curCanvas->changeCanvasMode(DRAW);
            annoContainer.setSelected(-1);
            if (curCanvas->getTaskMode() == SEGMENTATION || curCanvas->getTaskMode() == SEGMENTATION3D){
                ui->annotationButtonMoveUp->setEnabled(false);
                ui->annotationButtonMoveDown->setEnabled(false);
            }
        }else{                  // 进入SELECT模式，可编辑
            curCanvas->changeCanvasMode(SELECT);
            annoContainer.setSelected(ui->annoListWidget->row(items[0]));
            if (curCanvas->getTaskMode() == SEGMENTATION || curCanvas->getTaskMode() == SEGMENTATION3D){
                ui->annotationButtonMoveUp->setEnabled(true);
                ui->annotationButtonMoveDown->setEnabled(true);
            }
            if (curCanvas->getTaskMode() == SEGMENTATION3D){
                auto item = Seg3DAnnotationItem::castPointer(annoContainer.getSelectedItem());
                if (item!=nullptr){
                    int z = item->getStrokes()[0].z;
                    switchFile(z);
                }
            }
        }
        //! 可能是多余的？
        // if (curCanvas == canvas3d && canvas3d->getTaskMode() == SEGMENTATION3D)
        //     canvas3d->repaintSegAnnotation();
    });

    // 点击一个标注，如果是在3D检测下，则将focus设置到该标注的cube的中央
    connect(ui->annoListWidget, &QListWidget::itemClicked, [this](QListWidgetItem *_item){
        if (curCanvas == canvas3d && canvas3d->getTaskMode() == DETECTION3D){
            int row = ui->annoListWidget->row(_item);
            auto item = CubeAnnotationItem::castPointer(annoContainer.at(row));
            canvas3d->setFocusPos(item->getCube().center());
            switchFile(item->getCube().center().z);
        }
    });

    // annotation的右键菜单: delete
    ui->annoListWidget->setContextMenuPolicy(Qt::CustomContextMenu);
    connect(ui->annoListWidget, &QListWidget::customContextMenuRequested,
            this, &MainWindow::provideAnnoContextMenu);

    // undo & redo
    ui->actionUndo->setEnabled(false);
    connect(&annoContainer, &AnnotationContainer::UndoEnableChanged,
            ui->actionUndo, &QAction::setEnabled);
    connect(ui->actionUndo, &QAction::triggered, &annoContainer, &AnnotationContainer::undo);
    ui->actionRedo->setEnabled(false);
    connect(&annoContainer, &AnnotationContainer::RedoEnableChanged,
            ui->actionRedo, &QAction::setEnabled);
    connect(ui->actionRedo, &QAction::triggered, &annoContainer, &AnnotationContainer::redo);

    // 来自canvas的关于新标注的请求
    connect(canvas2d, &Canvas2D::newRectangleAnnotated, this, &MainWindow::getNewRect);
    connect(canvas2d, &Canvas2D::newStrokesAnnotated, this, &MainWindow::getNewStrokes);
    connect(canvas3d, &Canvas3D::newCubeAnnotated, this, &MainWindow::getNewCube);
    connect(canvas3d, &Canvas3D::newStrokes3DAnnotated, this, &MainWindow::getNewStrokes3D);
    // 来自canvas的关于修改或删除矩形（立方体）的标注的请求
    connect(canvas2d, &Canvas2D::removeRectRequest, &annoContainer, &AnnotationContainer::remove);
    connect(canvas2d, &Canvas2D::modifySelectedRectRequest, [this](int idx, QRect rect){
        std::shared_ptr<RectAnnotationItem> item =
                std::make_shared<RectAnnotationItem>(rect, annoContainer.getSelectedItem()->getLabel(),
                                                     annoContainer.getSelectedItem()->getId());
        annoContainer.modify(idx, std::static_pointer_cast<AnnotationItem>(item));
    });
    connect(canvas3d, &Canvas3D::removeCubeRequest, &annoContainer, &AnnotationContainer::remove);
    connect(canvas3d, &Canvas3D::modifySelectedCubeRequest, [this](int idx, Cuboid cube){
        std::shared_ptr<CubeAnnotationItem> item =
                std::make_shared<CubeAnnotationItem>(cube, annoContainer.getSelectedItem()->getLabel(),
                                                     annoContainer.getSelectedItem()->getId());
        annoContainer.modify(idx, std::static_pointer_cast<AnnotationItem>(item));
    });

    // anno changed -> canvas update
    connect(&annoContainer, &AnnotationContainer::annoChanged, this, &MainWindow::canvasUpdate);
    connect(&annoContainer, &AnnotationContainer::selectedChanged, this, &MainWindow::canvasUpdate);

    // anno changed -> ui list change
    connect(&labelManager, &LabelManager::colorChanged, [this](QString label, QColor color){
        for (int i=0;i<ui->annoListWidget->count();i++){
            auto item = ui->annoListWidget->item(i);
            if (item->text().split(' ')[0]==label)
                ui->annoListWidget->changeIconColorByIdx(i, color);
        }
    });
    connect(&annoContainer, &AnnotationContainer::AnnotationAdded,[this](const AnnoItemPtr &item){
        ui->annoListWidget->addCustomItemUncheckable(item->toStr(), labelManager.getColor(item->getLabel()));
    });
    connect(&annoContainer, &AnnotationContainer::AnnotationInserted,[this](const AnnoItemPtr &item, int idx){
        ui->annoListWidget->insertCustomItemUncheckable(item->toStr(), labelManager.getColor(item->getLabel()),idx);
    });
    connect(&annoContainer, &AnnotationContainer::AnnotationModified,[this](const AnnoItemPtr &item, int idx){
        ui->annoListWidget->changeTextByIdx(idx, item->toStr());
    });
    connect(&annoContainer, &AnnotationContainer::AnnotationRemoved,[this](int idx){
        ui->annoListWidget->removeCustomItemByIdx(idx);
    });
    connect(&annoContainer, &AnnotationContainer::allCleared,
            ui->annoListWidget, &QListWidget::clear);
    connect(&annoContainer, &AnnotationContainer::AnnotationSwap, [this](int idx){
        int selectedIdx = annoContainer.getSelectedIdx();

        auto item = ui->annoListWidget->takeItem(idx);
        ui->annoListWidget->insertItem(idx+1, item);

        if (selectedIdx == idx)
            ui->annoListWidget->item(idx)->setSelected(true);
        else if (selectedIdx == idx+1)
            ui->annoListWidget->item(idx+1)->setSelected(true);
    });
}


void MainWindow::_setupFileManager()
{
    // labelChangede,annoChanged -> setChangeNotSaved
    connect(&labelManager, &LabelManager::labelChanged, &fileManager, &FileManager::setChangeNotSaved);
    connect(&annoContainer, &AnnotationContainer::annoChanged, &fileManager, &FileManager::setChangeNotSaved);

    // prev/next action enable
    connect(&fileManager, &FileManager::prevEnableChanged, ui->actionPrevious_Image, &QAction::setEnabled);
    connect(&fileManager, &FileManager::nextEnableChanged, ui->actionNext_Image, &QAction::setEnabled);

    // fileListSetup -> update ui list
    connect(&fileManager, &FileManager::fileListSetup, [this](){
        ui->fileListWidget->clear();
        if (fileManager.getStatus() == Close) return;
        for (const QString &image: fileManager.allImageFiles()){
            ui->fileListWidget->addItem(FileManager::getNameWithExtension(image));
        }
        ui->fileListWidget->item(fileManager.getCurIdx())->setSelected(true);
    });

    // click file item to switch image
    connect(ui->fileListWidget, &QListWidget::itemSelectionChanged, [this](){
        auto items = ui->fileListWidget->selectedItems();
        if (items.length()==1){
            int idx = ui->fileListWidget->row(items[0]);
            if (idx != fileManager.getCurIdx()){
                if (curCanvas==canvas3d){
                    switchFile(idx);
                }else{
                    if (!switchFile(idx)){
                        ui->fileListWidget->item(fileManager.getCurIdx())->setSelected(true); // redundent?
                    }
                }
            }
            // listwidget的滚轮跟随移动
            if (ui->fileListWidget->count()>1){
                int row = ui->fileListWidget->row(items[0]);
                auto scroll = ui->fileListWidget->verticalScrollBar();
                scroll->setValue(scroll->minimum()+
                                  (scroll->maximum()-scroll->minimum())*row/(ui->fileListWidget->count()-1));
            }
        }
    });

    // 3D情况下，focus的移动影响到图像的切换
    connect(canvas3d, &Canvas3D::focus3dMoved, [this](Point3D focus){
        fileManager.selectFile(focus.z);
        ui->fileListWidget->item(focus.z)->setSelected(true);
    });
}

bool MainWindow::switchFile(int idx)
{
    if (curCanvas == canvas2d){
        if (!_checkUnsaved()) return false;

        labelManager.allClear();
        annoContainer.allClear();

        fileManager.selectFile(idx);
        _loadJsonFile(fileManager.getLabelFile());
        _loadJsonFile(fileManager.getCurrentOutputFile());
        fileManager.resetChangeNotSaved();

        canvas2d->close();
        canvas2d->loadPixmap(fileManager.getCurrentImageFile());
        adjustFitWindow();

        ui->fileListWidget->item(idx)->setSelected(true);
        return true;
    }else if (curCanvas == canvas3d){
        fileManager.selectFile(idx);
        ui->fileListWidget->item(idx)->setSelected(true);
        Point3D focusPos = canvas3d->getFocusPos();
        canvas3d->setFocusPos(Point3D(focusPos.x, focusPos.y, idx));
        return true;
    }else{
        throw "abnormal curCanvas";
    }
}

void MainWindow::on_actionSave_triggered()
{
    if (curCanvas == canvas2d){ // 2D mode
            if (fileManager.getMode()==SingleImage){
                QJsonObject json;
                json.insert("labels", labelManager.toJsonArray());
                json.insert("annotations", annoContainer.toJsonArray());
                FileManager::saveJson(json, fileManager.getCurrentOutputFile());

                if (canvas2d->getTaskMode()==SEGMENTATION)
                    _saveSegmentImageResults();
            }else if (fileManager.getMode()==MultiImage){
                QJsonObject labelJson;
                labelJson.insert("labels", labelManager.toJsonArray());
                FileManager::saveJson(labelJson, fileManager.getLabelFile());
                QJsonObject annoJson;
                annoJson.insert("annotations", annoContainer.toJsonArray());
                FileManager::saveJson(annoJson, fileManager.getCurrentOutputFile());

                if (canvas2d->getTaskMode()==SEGMENTATION)
                    _saveSegmentImageResults();
            }
        }else if (curCanvas == canvas3d){ // 3D mode
            if (canvas3d->getTaskMode() == DETECTION3D){
                QJsonObject json;
                json.insert("labels", labelManager.toJsonArray());
                json.insert("annotations", annoContainer.toJsonArray());
                FileManager::saveJson(json, fileManager.getCurrentOutputFile());
            }else if (canvas3d->getTaskMode() == SEGMENTATION3D){
                QJsonObject json;
                json.insert("labels", labelManager.toJsonArray());
                json.insert("annotations", annoContainer.toJsonArray());
                FileManager::saveJson(json, fileManager.getCurrentOutputFile());

                if (canvas3d->getTaskMode()==SEGMENTATION3D){
                    _saveSegment3dImageResults();
                }
            }
        }
        fileManager.resetChangeNotSaved();
}

void MainWindow::unableFileActions()
{
    for (auto action: fileRelatedActions)
        action->setEnabled(false);
}

void MainWindow::_loadJsonFile(QString fileName)
{
    QFileInfo checkFile = QFileInfo(fileName);
    if (checkFile.exists() && checkFile.isFile()){
        try {
            QJsonObject json = FileManager::readJson(fileName);
            labelManager.fromJsonObject(json);
            annoContainer.fromJsonObject(json, curCanvas->getTaskMode());
        } catch (FileException &e) {
            QMessageBox::warning(this, "File Error", e.what());
        } catch (JsonException &e) {
            QString msg;
            msg = QString("The saved json file is broken.\n")
                    +"Error message: "+e.what()+"\n"
                    +"Please check or delete the json file.\n";
            QMessageBox::warning(this, "Json Error", msg);
        }
    }
}

bool MainWindow::_checkUnsaved()
{
    if (fileManager.hasChangeNotSaved()){
            if (ui->actionSave_Automatically->isChecked())
                on_actionSave_triggered();
            else{
                int ret = QMessageBox::warning(this, QObject::tr("Warning"),
                                               QObject:: tr("The document has been modified.\n"
                                                            "Do you want to save your changes?\n"
                                                            "Note: you can check AutoSave option in the menu.\n"),
                                               QMessageBox::Save | QMessageBox::Discard | QMessageBox::Cancel,
                                               QMessageBox::Save);
                switch (ret) {
                case QMessageBox::Save:
                    on_actionSave_triggered();
                    break;
                case QMessageBox::Discard:
                    break;
                case QMessageBox::Cancel:
                    return false;
                default:
                    break;
                }
            }
        }
        return true;
}
